#include "planner/optimizer.h"

#include "common/defer.h"
#include "common/formator.h"
#include "common/scoped_timer.h"
#include "planner/normalizer.h"
#include "planner/optimization_context.h"
#include "planner/optimizer_task_pool.h"
#include "planner/property_enforcer.h"
#include "planner/rule.h"

extern "C" {
#include <utils/elog.h>
}

namespace optimizer {

group_id_t Optimizer::BuildPlanTree(std::shared_ptr<Operator> op_tree, PropertySet *required_props) {
  GroupExpression *gexpr = nullptr;

  // normalize the operator tree
  // 1. push filter down
  // 2. push projection down
  // 3. simple expression
  // 4. subquery unnesting
  // 5. drive output
  // 6. drive target
  {
    ScopedDefer defer{[op_tree] { LOG_PRINTLN("before normalization:\n%s", Format::ToString(op_tree).c_str()); },
                      [op_tree] { LOG_PRINTLN("after normalization:\n%s", Format::ToString(op_tree).c_str()); }};
    Normalizer normalizer(op_tree);
    op_tree = normalizer.Normalize();
  }

  context_->RecordOptimizerNodeIntoGroup(std::move(op_tree), &gexpr);

  group_id_t root_id = gexpr->GetGroupID();

  try {
    OptimizeLoop(root_id, required_props);
    LOG_PRINTLN("After optimization for Memo:\n%s", context_->GetMemo().ToString().c_str());
    return root_id;
  } catch (std::exception &e) {
    LOG_PRINTLN("Optimize Loop ended prematurely: %s", e.what());
    return -1;
  }
}

void Optimizer::OptimizeLoop(group_id_t root_group_id, PropertySet *required_props) {
  auto *root_context = new OptimizationContext(context_, required_props->Copy());
  context_->AddOptimizationContext(root_context);

  // Perform rewrite first
  context_->PushTask(new TopDownRewrite(root_group_id, root_context, RuleSetName::PREDICATE_PUSH_DOWN));
  context_->PushTask(new BottomUpRewrite(root_group_id, root_context, RuleSetName::UNNEST_SUBQUERY, false));
  ExecuteTaskStack(root_group_id, root_context);

  // Perform optimization after the rewrite
  Memo &memo = context_->GetMemo();
  LOG_PRINTLN("Starting optimization for Memo:\n%s", memo.ToString().c_str());
  context_->PushTask(new OptimizeGroup(memo.GetGroupByID(root_group_id), root_context));

  // Derive stats for the only one logical expression before optimizing
  context_->PushTask(new DeriveStats(memo.GetGroupByID(root_group_id)->GetLogicalExpression(), root_context));

  ExecuteTaskStack(root_group_id, root_context);
}

void Optimizer::ExecuteTaskStack(group_id_t root_group_id, OptimizationContext *root_context) const {
  auto *root_group = context_->GetMemo().GetGroupByID(root_group_id);
  const auto &required_props = root_context->GetRequiredProperties();

  uint64_t elapsed_time = 0;

  auto &task_stack = context_->GetTaskPool();

  // Iterate through the task stack
  while (!task_stack.Empty()) {
    // Check to see if we have at least one plan, and if we have exceeded our
    // timeout limit
    if (elapsed_time >= task_execution_timeout_ && root_group->HasExpressions(required_props)) {
      throw std::runtime_error("Optimizer task execution timed out");
    }

    // CHECK_FOR_INTERRUPTS();

    uint64_t task_runtime = 0;
    auto *task = task_stack.Pop();
    {
      ScopedTimer<std::chrono::milliseconds> timer(&task_runtime);
      task->Execute();
    }
    delete task;
    elapsed_time += task_runtime;
  }
}

}  // namespace optimizer
